import argparse
import pandas as pd
import clevercsv
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="生成CSV文件中指定字段的相关性矩阵图")
    parser.add_argument("input", type=str, help="输入的CSV文件路径")
    parser.add_argument(
        "-f", "--fields", 
        nargs="+", 
        default=['vchg','vbat','frontshell_therm','board_up_therm','board_down_therm','air_inlet_therm','backshell_therm','sxr_therm','wlan_pa_therm','usb_therm','bat_therm','cpu-0-0-0','gpuss-0','nspss-0','Fan1_PWM','Fan1_state','Fan2_PWM','Fan2_state','VST1_T_sensor','VST2_T_sensor'],
        help="需要计算相关性的字段名（空格分隔）"
    )
    parser.add_argument(
        "-t", "--target_field", 
        type=str, 
        default='Y_std', 
        help="目标字段（只计算目标字段与其他字段的相关性）"
    )
    parser.add_argument(
        "-o", "--output", 
        type=str, 
        default=None, 
        help="输出图片路径（未指定时直接显示）"
    )
    return parser.parse_args()

def validate_fields(df, fields, target_field=None):
    """检查字段是否存在于数据中"""
    missing_fields = [f for f in fields if f not in df.columns]
    if missing_fields:
        raise ValueError(f"字段不存在于CSV中: {missing_fields}")
    

def clean_data(df, fields, target_field=None):
    """
    清洗数据：将百分比字符串转换为小数，处理逗号分隔符
    例如: '44%' -> 0.44, '1,000' -> 1000.0
    """
    # 创建需要清洗的字段列表（包括目标字段）
    clean_fields = fields.copy()
    if target_field:
        clean_fields.append(target_field)
    
    for field in clean_fields:
        # 检查列是否是对象类型（通常是字符串）
        if df[field].dtype == 'object':
            # 移除空格和逗号，并处理百分比
            df[field] = df[field].astype(str).str.replace(' ', '').str.replace(',', '')
            # 检测并转换百分比值
            if df[field].str.contains('%').any():
                df[field] = df[field].str.rstrip('%').astype(float) / 100.0
            else:
                # 尝试转换为数值类型
                try:
                    df[field] = pd.to_numeric(df[field], errors='raise')
                except ValueError:
                    raise ValueError(f"字段 '{field}' 包含无法转换的非数值数据")
    
    return df

def plot_correlation_matrix(df, fields, target_field=None, output_path=None):
    """生成并保存/显示相关性矩阵热图"""
    # 计算相关性
    if target_field and len(target_field) > 1:
        # 只计算目标字段与其他字段的相关性
        corr_series = df[fields].corrwith(df[target_field])
        
        # 按相关性值从大到小排序
        corr_sorted = corr_series.sort_values(ascending=False)
        corr_df = corr_sorted.to_frame()
        corr_df.columns = [f'Corr with {target_field}']
        
        # 绘制单列热图（按相关性排序）
        plt.figure(figsize=(8, max(6, len(fields) * 0.6)))
        ax = sns.heatmap(
            corr_df, 
            annot=True, 
            cmap="coolwarm", 
            vmin=-1, 
            vmax=1, 
            fmt=".3f",
            linewidths=0.5,
            cbar_kws={"shrink": 0.8}
        )
        plt.title(f"Correlation with {target_field} (Sorted)", fontsize=14)
        
        # 添加排序指示箭头
        ax.text(0.5, -0.15, "↑ Higher Correlation", 
                ha='center', va='center', transform=ax.transAxes, color='darkred')
        ax.text(0.5, -0.2, "↓ Lower Correlation", 
                ha='center', va='center', transform=ax.transAxes, color='darkblue')
    else:
        # 计算所有字段之间的相关性矩阵
        corr = df[fields].corr()
        
        # 创建遮罩（仅显示下三角）
        mask = np.zeros_like(corr, dtype=bool)
        mask[np.triu_indices_from(mask)] = True

        # 绘制完整矩阵热图
        plt.figure(figsize=(10, 8))
        sns.heatmap(
            corr, 
            annot=True, 
            cmap="coolwarm", 
            vmin=-1, 
            vmax=1, 
            fmt=".2f",
            linewidths=0.5,
            mask=mask
        )
        # 设置标题和标签
        plt.title("Correlation Matrix", fontsize=14)
        
        # 旋转x轴标签并调整位置
        plt.xticks(rotation=45, ha='right', rotation_mode='anchor')
        
        # 设置y轴标签位置
        plt.yticks(va="center")
    
    # 自动调整布局
    plt.tight_layout()
    
    # 处理输出
    if output_path:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_path, bbox_inches="tight", dpi=300)
        print(f"图片已保存至: {output_path}")
    else:
        plt.show()

def main():
    args = parse_args()
    if args.target_field and args.target_field.lower() == 'none':
        args.target_field = None
    # 读取CSV文件
    df = clevercsv.read_dataframe(args.input)
    
    # 检查字段有效性
    validate_fields(df, args.fields, args.target_field)
    
    # 数据清洗：转换百分比和数值格式
    df = clean_data(df, args.fields, args.target_field)
    
    # 生成并保存/显示热图
    plot_correlation_matrix(df, args.fields, args.target_field, args.output)

if __name__ == "__main__":
    main()